import datetime
import logging
import json
import os
import re
import subprocess
import tempfile
import time
import toml
import traceback

import pytz
from absl import app, flags

from coin.base.hostname import get_hostname
from coin.base.param_util import to_list, split_string
from coin.exchange.base.order_gateway_logger import OrderGatewayLogProvider
from coin.proto.coin_request_pb2 import AccountRequestProto
from coin.support.accounting.logic.key_util import get_key_file
from xunkemgmt_client.client.util.query_util import query_accounts


FLAGS = flags.FLAGS


def _load_proxy_config(proxy_config_file):
  proxy = None
  if proxy_config_file is not None:
    assert os.path.exists(proxy_config_file), proxy_config_file
    with open(proxy_config_file) as f:
      proxy = json.loads(f.read())
  return proxy


def _get_proxy_config_file(proxy_info_file, account_request):
  proxy_info = toml.load(proxy_info_file)
  me = f'{account_request.market_type}.{account_request.exchange}'
  chosen = proxy_info.get(me, [])
  ret = None
  for item in to_list(chosen):
    assert ret is None, ret
    if 'regex' not in item or re.match(item['regex'], account_request.owner) is not None:
      ret = item
      break
  ret = ret if ret is None else \
        os.path.join(os.path.dirname(proxy_info_file), ret['proxy'])
  return ret


def _validate_account_request(account_request):
  assert account_request.market_type is not None
  assert account_request.exchange is not None
  assert account_request.owner is not None
  assert account_request.api_version is not None


def _get_key_root_path(account_request):
  if account_request.business_unit is not None:
    if account_request.business_unit in ('DigitalFund'):
      return '../../fund_key/view'
    else:
      return '../../coin_key/view'
  else:
    return None


def _get_product_info_file(pi_prefix, pi_mea):
  current_date = datetime.datetime.utcnow()
  assert os.path.exists(pi_prefix), pi_prefix
  pi_file = os.path.join(pi_prefix, current_date.strftime("%Y%m%d"), f'{pi_mea}.json')
  if os.path.exists(pi_file):
    return pi_file
  return os.path.join('data/coin2/product_info', f'{pi_mea}.json')


def generate_account_loggers(request_list, log_root=None, machine=None):
  assert log_root is not None
  machine = machine or get_hostname()
  account_loggers = {}
  for request in request_list:
    account_request_proto = AccountRequestProto(
        market_type=request.market_type,
        exchange=request.exchange,
        owner=request.owner)
    account_logger = OrderGatewayLogProvider(
        account_request_proto, log_root, machine)
    account_key = (request.market_type, request.exchange, request.owner)
    assert account_key not in account_loggers
    account_loggers[account_key] = account_logger
  return account_loggers


def account_request_to_driver_dict(account_request, proxy_config_file):
  _validate_account_request(account_request)
  market_type = account_request.market_type
  exchange = account_request.exchange
  owner = account_request.owner
  api_version = account_request.api_version
  mea = f'{market_type}.{exchange}.{api_version}'
  if market_type in ('ExWallet', 'Margin'):
    pi_mea = f'Spot.{exchange}.{api_version}'
  else:
    pi_mea = mea
  pi_file = _get_product_info_file('../../coin_product_info/data', pi_mea)
  with open(pi_file) as f:
    pi = json.load(f)
    assert pi_mea == pi['mea'], (pi_mea, pi['mea'])
    norms = [product['symbol'] for product in pi['product_infos']]

  key_filepath = get_key_file(
    market_type, exchange, owner, api_version,
    "view", key_root=_get_key_root_path(account_request)
  )
  products = {'norms': norms}
  connection_config = {'key_filepath': key_filepath, 'use_cross_mode': True}
  if proxy_config_file is not None:
    proxy = _load_proxy_config(proxy_config_file)
    connection_config['http_proxy'] = proxy['http']
  else:
    if account_request.exchange == 'Bithumb':
      connection_config['rest_host'] = 'api.bithumb.com'
    if account_request.exchange == 'Prex':
      connection_config['rest_host'] = 'api.aqx.com'
      connection_config["ws_disable_subscribe"] = True
  portfolio_margin_trading = 'portfolio' in owner.lower() and exchange == 'Binance'
  rest_query_config = {'query_account_balance_period_sec': 30,
                       'query_account_position_period_sec': 30,
                       'query_open_orders_period_sec': 100000000,
                       'query_fill_period_sec': 100000000,
                       'query_cancel_orders_period_sec': 100000000}
  order_logger_config = {
      'store_parsed_message': False,
      'as_account_logger': True
  }
  order_gateway_config = {
    'log_private_exchange_message': False
  }
  order_config = {
      'connection_config': connection_config,
      'mea': mea,
      'order_gateway_config': order_gateway_config,
      'order_logger_config': order_logger_config,
      'portfolio_margin_trading': portfolio_margin_trading,
      'products': products,
      'rest_query_config': rest_query_config,
      'type': 'ORDER_SYSTEM_TYPE_LIVE_OG'
  }
  driver = {'order': {'exchanges': {exchange: order_config}}}
  return driver


def collect_account_info(request_list, proxy_info_file):
  if len(request_list) == 0:
    return
  os.makedirs(FLAGS.log_dir, exist_ok=True)
  collector_procs = []
  for request in request_list:
    try:
      proxy_config_file = _get_proxy_config_file(proxy_info_file, request) \
                          if proxy_info_file is not None else None
      driver_dict = account_request_to_driver_dict(request, proxy_config_file)
      acct_key = f'{request.market_type}.{request.exchange}.{request.owner}'
      dt_str = datetime.datetime.utcnow().replace(tzinfo=pytz.UTC).strftime(
            '%Y%m%d-%H%M%S%z')
      log_file = os.path.join(
          FLAGS.log_dir,
          f'account_info_collector.{acct_key}.{dt_str}')
      out_driver_json = tempfile.mkstemp(
          dir='/tmp', prefix=f'{acct_key}.driver_config.', suffix='.json')[-1]
      with open(out_driver_json, 'w') as fwrite:
        json.dump(driver_dict, fwrite, indent=2, sort_keys=True)
    except Exception as e:
      logging.error(traceback.format_exc())
      continue
    cmd = 'source ${HOME}/miniconda3/etc/profile.d/conda.sh' + ' && ' + \
          'conda activate %s' % FLAGS.env + ' && ' + \
          'export GLOG_logtostderr=1' + ' && ' + \
          FLAGS.binary_path + ' ' + \
          '--driver=%s --exit_after_sec %s > %s 2>&1' % (
              out_driver_json, FLAGS.duration, log_file)
    collector_procs.append(subprocess.Popen(cmd, shell=True))
  start_time = time.time()
  while True:
    subprocess_exit = all([proc.poll() is not None for proc in collector_procs])
    if subprocess_exit:
      return
    if time.time() - start_time > FLAGS.duration + 60:
      logging.info('kill subprocess due to timeout.')
      for proc in collector_procs:
        proc.terminate()
    time.sleep(10)


def main(_):
  logging.info('start account info collection')
  request_list = query_accounts(
      active=True,
      market_types=split_string(FLAGS.market_type),
      exchanges=split_string(FLAGS.exchange),
      owners=split_string(FLAGS.owner),
      as_proto=True)
  collect_account_info(request_list, FLAGS.proxy_info_file)


if __name__ == '__main__':
  logging.basicConfig(level='INFO', format='%(levelname)8s %(asctime)s %(message)s')

  flags.DEFINE_string('market_type', None, 'Comma separated market_type.')

  flags.DEFINE_string('exchange', None, 'Comma separated exchange.')

  flags.DEFINE_string('owner', None, 'Comma separated account owner.')

  flags.DEFINE_string('proxy_info_file', None, 'proxy_info_file')

  flags.DEFINE_string('env', 'coin2_motion_env_migration', 'conda enviroment')

  flags.DEFINE_string('binary_path',
                      '../../coin_binary/coin2-bazel-bin-3.10/cc/appcoin2/support/order/account_info_collector',
                      'collector binary path')

  flags.DEFINE_integer(
      'duration',
      3600,
      'Duration for running collector in seconds. Eg. If set to 300, program will '
      'stop after running 300 seconds. If set to None, program will run forever')

  app.run(main)
